from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import HTMLResponse, StreamingResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), "lib", "bin"))
import hfslib
from io import BytesIO
from PIL import Image
import numpy as np
import aiofiles
import uvicorn

SERVER_ROOT = ''

icons = {}

app = FastAPI()


app.mount('/static', StaticFiles(directory='statics'))

@app.get('/', response_class=HTMLResponse)
async def index():
    async with aiofiles.open("public/index.html", 'rb') as f:
        return HTMLResponse(await f.read())

class FileInfo(BaseModel):
    name: str
    size: int | None
    time: int | float
    path: str
    is_dir: bool

@app.get('/listdir', response_model=list[FileInfo])
def listdir(root: str | None = None):
    root = SERVER_ROOT if root is None else os.path.join(SERVER_ROOT, root)
    return [
        (lambda fn,path: FileInfo(
            name=fn,
            size=None if os.path.isdir(path) else os.stat(path).st_size,
            time=os.stat(path).st_mtime,
            is_dir=os.path.isdir(path),
            path=os.path.relpath(path, SERVER_ROOT)
        ))(fn, os.path.join(root, fn))
        for fn in os.listdir(root)
    ]

@app.get('/icon', response_class=StreamingResponse)
def geticon(path: str):
    if path not in icons:
        r = hfslib.get_icon(os.path.join(SERVER_ROOT, path))
        img = Image.fromarray(np.array(r).astype('uint8')).convert('RGBA')
        bs = BytesIO()
        img.save(bs, format='png')
        bs = bs.getvalue()
        img.close()
        icons[path] = bs
    return StreamingResponse(BytesIO(icons[path]), media_type='image/png')

@app.get('/download', response_class=FileResponse)
async def download(path: str):
    return FileResponse(os.path.join(SERVER_ROOT, path))

@app.post('/upload')
async def upload(file: UploadFile = File(), root: str = Form(default='')):
    root = os.path.join(SERVER_ROOT, root)
    async with aiofiles.open(os.path.join(root, file.filename), 'wb') as f:
        while True:
            t = await file.read(2**26)
            if len(t) < 1:
                break
            await f.write(t)
    return {'status': 'ok'}

if __name__=='__main__':
    if len(sys.argv) == 2:
        SERVER_ROOT = sys.argv[1]
        uvicorn.run(app, port = 9999)